Evaluate YOLOv5 Detector Performance

Training losses and performance metrics are saved to Tensorboard and also to a logfile.

# Start tensorboard
# Launch after you have started training
# logs save in the folder "runs"
#%load_ext tensorboard
%reload_ext tensorboard
%tensorboard --logdir /content/yolov5/runs

Run Inference With Trained Weights

Run inference with the best pretrained checkpoint on the test set.

!cp /content/drive/MyDrive/main_image_star-forming_region_carina_nircam_final-5mb.jpg /content/datasets/JWST-stars-1/test/images

!python detect.py --weights /content/yolov5/runs/train/exp/weights/best.pt --img 416 --conf 0.1 --source {dataset.location}/test/images --hide-labels --hide-conf

detect: weights=['/content/yolov5/runs/train/exp/weights/best.pt'], source=/content/datasets/JWST-stars-1/test/images, data=data/coco128.yaml, imgsz=[416, 416], conf_thres=0.1, iou_thres=0.45, max_det=1000, device=, view_img=False, save_txt=False, save_conf=False, save_crop=False, nosave=False, classes=None, agnostic_nms=False, augment=False, visualize=False, update=False, project=runs/detect, name=exp, exist_ok=False, line_thickness=3, hide_labels=True, hide_conf=True, half=False, dnn=False
YOLOv5 🚀 v6.1-295-gf8722b4 Python-3.7.13 torch-1.11.0+cu102 CUDA:0 (Tesla T4, 15110MiB)

Fusing layers... 
Model summary: 213 layers, 7012822 parameters, 0 gradients, 15.8 GFLOPs
image 1/4 /content/datasets/JWST-stars-1/test/images/000_png.rf.4ae0a50400036394796e15cbbc57ca5e.jpg: 416x416 19 stars, Done. (0.009s)
image 2/4 /content/datasets/JWST-stars-1/test/images/003_png.rf.16499f5a19b37fe576efb1d5121f51fd.jpg: 416x416 9 stars, Done. (0.009s)
image 3/4 /content/datasets/JWST-stars-1/test/images/021_png.rf.1ffca634cbafa2d72a8861c0f062c469.jpg: 416x416 33 stars, Done. (0.009s)
image 4/4 /content/datasets/JWST-stars-1/test/images/main_image_star-forming_region_carina_nircam_final-5mb.jpg: 256x416 21 stars, Done. (0.007s)
Speed: 0.3ms pre-process, 8.4ms inference, 1.0ms NMS per image at shape (1, 3, 416, 416)
Results saved to runs/detect/exp

#display inference on ALL test images

import glob
from IPython.display import Image, display

for imageName in glob.glob('/content/yolov5/runs/detect/exp/*.jpg'): #assuming JPG
    display(Image(filename=imageName))
    print("\n")

Future Works

This study presents a series of future improvements:

• I employed the smallest, simpler model of the family, YOLOv5n, with only 1.9 million parameters. Other more complex and better models are available, up to the YOLOv5x, which contains 86.7 million parameters;
• I did not use the raw full quality images, which for sure could improve the performance of the model, but also increase substantially the training time;
• My training set consisted of only 54 images, extracted from the full image. The training and testing sets can be better worked on, and also with more augmentations;
• The initial weights used for the model are the weights on the generic COCO pre-trained checkpoint. These may not the best initial weights for stars detection;
• Multiple classes are possible, for different celestial bodies classification, and not only star detection;
• Other images from the JWST can also be applied to the network.

Overall, even though the YOLOv5 model was not designed to detect stars, it can be adapted and trained for such performance. Here I just applied a simple computational vision approach on the James Webb Space Telescope images, which can be improved much more.